import pymysql
from database_manager import get_db_connection

def update_stock_data():
    """
    更新数据库中的股票数据：
    1. 为null的industry字段补充行业信息
    2. 修改symbol_standardized字段，去掉前两个字符
    """
    conn = get_db_connection()
    if not conn:
        print("无法连接数据库")
        return
    
    try:
        with conn.cursor() as cursor:
            # 1. 补充行业信息
            # 创建行业映射字典，基于股票代码前缀
            industry_mapping = {
                '600': '银行', '601': '银行', '603': '银行',
                '000': '房地产', '001': '房地产',
                '002': '制造业',
                '300': '科技', '301': '科技',
                '688': '科创板',
                # 其他行业映射可以根据需要添加
            }
            
            # 更新industry为null的记录，根据股票代码前缀设置行业
            cursor.execute("SELECT id, symbol_standardized FROM stock_info WHERE industry IS NULL")
            null_industry_stocks = cursor.fetchall()
            
            updated_count = 0
            for stock in null_industry_stocks:
                # 提取股票代码的前三位作为行业判断依据
                prefix = stock['symbol_standardized'][:3]
                industry = industry_mapping.get(prefix, '其他')
                
                cursor.execute(
                    "UPDATE stock_info SET industry = %s WHERE id = %s",
                    (industry, stock['id'])
                )
                updated_count += cursor.rowcount
            
            print(f"已更新 {updated_count} 条股票的行业信息")
            
            # 2. 修改symbol_standardized字段，去掉前两个字符
            # 先获取所有记录
            cursor.execute("SELECT id, symbol_standardized FROM stock_info")
            all_stocks = cursor.fetchall()
            
            modified_count = 0
            for stock in all_stocks:
                current_symbol = stock['symbol_standardized']
                # 检查是否需要修改（长度大于2）
                if len(current_symbol) > 2:
                    new_symbol = current_symbol[2:]  # 去掉前两个字符
                    cursor.execute(
                        "UPDATE stock_info SET symbol_standardized = %s WHERE id = %s",
                        (new_symbol, stock['id'])
                    )
                    modified_count += cursor.rowcount
            
            print(f"已修改 {modified_count} 条股票的标准化代码")
            
            # 提交事务
            conn.commit()
            print("数据库更新完成")
            
    except Exception as e:
        print(f"更新数据库时出错: {e}")
        conn.rollback()
    finally:
        conn.close()

if __name__ == "__main__":
    update_stock_data()
